import sys

# 添加自定义的 OpenCV 编译路径（仅当前脚本会生效）
sys.path.insert(0, "/usr/local/lib/python3.8/site-packages/cv2/python-3.8")
import cv2


import os
import numpy as np
import torch
import logging

# 兼容旧代码对 np.bool 的引用
if not hasattr(np, "bool"):
    np.bool = bool  # 或者 np.bool_，视你是否需要 NumPy 标量类型

from torchvision import transforms
from typing import Sequence, Union
from copy import deepcopy
from collections import defaultdict
from mmengine.structures import PixelData
from mmseg.apis import inference_model, init_model

from .preprocess import (
    LoadImageFromNDArray,
    Resize,
    PackSegInputs,
    SegDataPreProcessor,
    resize,
    SegDataSample,
)
from .post_pipeline import mapillay_postprocess
from .metainfo_custom import metainfo
from .vis_utils import show_result_pyplot

ImageType = Union[str, np.ndarray, Sequence[str], Sequence[np.ndarray]]


logging.info(cv2.__file__)
logging.info(cv2.__version__)


class PytorchInfer:
    """
    image_seg_and_point
    该类用于图像分割和目标点生成
    """

    def __init__(self, checkpoint, infer_cfg):
        # cuda.init()
        self.model = init_model(infer_cfg["config_path"],
                           checkpoint,
                           device='cuda:0',
                           custom_palette=None
                           )
        self.use_postprocess = infer_cfg["use_postprocess"]
        # deeplabv3plus原始类别：
        # 0:"车行道路类", 1:"人行道路类", 2:"自然区域", 3:"禁行区",
        # 4:"交通信号灯", 5:"人行横道", 6:"道路标线", 7:"路沿石",
        # 8:"障碍物", 9:"背景",

        # 映射后：
        # 0：背景（背景9）
        # 1：平坦(车行道路0、人行道路1、人行横道5、道路标线6)，
        # 2：粗糙（自然区域2），
        # 3: 颠簸（路沿石7），
        # 4：禁区（禁行区3），
        # 5：障碍物（交通信号灯4、障碍物8），

        self.label_map = {0: 1, 1: 1, 2: 2, 3: 4, 4: 5, 5: 1, 6: 1, 7: 3, 8: 5, 9: 0}

    def initialize(self):
        pass

    def infer(self, frame: np.ndarray):

        result = inference_model(self.model, frame)

        return result

    def remap_labels(self, res):
        if isinstance(res, SegDataSample):
            new_res = deepcopy(res)
            seg_map = res.pred_sem_seg.data
            seg_map_new = torch.zeros_like(seg_map)
            for old_label, new_label in self.label_map.items():
                seg_map_new[seg_map == old_label] = new_label
            new_res.pred_sem_seg.data = seg_map_new
        elif isinstance(res, np.ndarray):
            new_res = np.zeros_like(res)
            for old_label, new_label in self.label_map.items():
                new_res[res == old_label] = new_label
        elif isinstance(res, torch.Tensor):
            new_res = torch.zeros_like(res)
            for old_label, new_label in self.label_map.items():
                new_res[res == old_label] = new_label
        else:
            raise ValueError

        return new_res

    def clean(self):
        pass

